package com.starhub.application.rag.util;

import java.util.List;
import java.util.Map;
import java.util.HashMap;
import java.util.stream.Collectors;
import java.util.Comparator;

import org.springframework.stereotype.Component;

import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.result.RetrieveResult;
import com.starhub.integration.milvus.langchain.MilvusMetadataFilterMapper;

import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import lombok.extern.slf4j.Slf4j;

/**
 * Milvus条件构建工具类
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-25
 */
@Slf4j
@Component
public class MilvusConditionBuilder {

    /**
     * 构建Milvus元数据过滤条件，最基础的条件模式，拼接所有item_id
     * 
     * @param mainIds
     * @return
     */
    public String buildNormalCondition(List<RetrieveResult> results) {
        if(results == null || results.isEmpty()) {
            return "";
        }

        String conditions = "";

        List<String> filters = results.stream().map(result -> {
            List<String> mainIds = result.getMainIds();
            if (!mainIds.isEmpty()) {
                Filter filter = new MetadataFilterBuilder(RagConstant.EmbeddingMetaData.ITEM_ID)
                        .isIn(mainIds);
                String condition = MilvusMetadataFilterMapper.map(filter, RagConstant.LangchainMilvusFields.FIELD_METADATA);
                return condition;
            }
            return null;
        }).collect(Collectors.toList());

        if(filters.isEmpty()) {
            return "";
        }

        // 使用 OR 连接多个条件
        conditions = String.join(" || ", filters.stream()
                .filter(f -> f != null && !f.isEmpty())
                .collect(Collectors.toList()));

        return conditions;
    }


    /**
     * 构建元数据关联检索条件
     * 将所有维度的关联相似度加权平均，rerank的依据
     * 合并所有RetrieveResult的mainIdsScore，获取相似度总和最高的10个key
     * 
     * @param results 检索结果列表
     * @return 构建的条件字符串
     */
    public String buildMetaLinkResult(List<RetrieveResult> results) {
        if (results == null || results.isEmpty()) {
            return "";
        }

        // 合并所有RetrieveResult的mainIdsScore
        Map<String, Float> mergedScores = new HashMap<>();
        
        for (RetrieveResult result : results) {
            Map<String, Float> mainIdsScore = result.getMainIdsScore();
            if (mainIdsScore != null && !mainIdsScore.isEmpty()) {
                for (Map.Entry<String, Float> entry : mainIdsScore.entrySet()) {
                    String key = entry.getKey();
                    Float score = entry.getValue();
                    
                    // 如果key已存在，累加相似度值
                    mergedScores.merge(key, score, Float::sum);
                }
            }
        }

        if (mergedScores.isEmpty()) {
            return "";
        }

        // 按相似度总和排序，获取最高的10个key
        List<String> topKeys = mergedScores.entrySet().stream()
                .sorted(Map.Entry.<String, Float>comparingByValue(Comparator.reverseOrder()))
                .limit(RagConstant.RETRIEVE_EMBEDDING_COUNT)
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());

        if (topKeys.isEmpty()) {
            return "";
        }

        // 构建过滤条件
        Filter filter = new MetadataFilterBuilder(RagConstant.EmbeddingMetaData.ITEM_ID)
                .isIn(topKeys);
        String condition = MilvusMetadataFilterMapper.map(filter, RagConstant.LangchainMilvusFields.FIELD_METADATA);

        return condition;
    }

    /**
     * 构建Milvus元数据过滤条件
     * 
     * @param results
     * @return
     */
    public String buildChatCondition(List<RetrieveResult> results) {
        return "";
    }
}
